package com.gouwudai.gobbs.test;

import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;

import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConfig;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.Column;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.DefaultDataSet;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.ITable;
import org.dbunit.dataset.ITableIterator;
import org.dbunit.dataset.ITableMetaData;
import org.dbunit.dataset.filter.IColumnFilter;
import org.dbunit.operation.AbstractOperation;
import org.dbunit.operation.CompositeOperation;
import org.dbunit.operation.DatabaseOperation;

/**
 * <pre>
 *  Title: 处理MSSQL IDENTITY_INSERT的InsertIdentityOperation
 *  Description: 修改了DBUNIT的InsertIdentityOperation实现，取消自动提交
 * </pre>
 * 
 * @author <a href="mailto:littcai@hotmail.com">空心大白菜</a>
 * @date 2007-4-13
 * @version 1.0
 * 
 * <pre>
 * 修改记录
 *    修改后版本:     修改人：  修改日期:     修改内容: 
 * </pre>
 */
public class InsertIdentityOperation extends AbstractOperation
{
    public static final String PROPERTY_IDENTITY_COLUMN_FILTER = "http://www.dbunit.org/properties/mssql/identityColumnFilter";

    public static final DatabaseOperation INSERT = new InsertIdentityOperation(
            DatabaseOperation.INSERT);

    public static final DatabaseOperation CLEAN_INSERT = new CompositeOperation(
            DatabaseOperation.DELETE_ALL, new InsertIdentityOperation(
                    DatabaseOperation.INSERT));

    public static final DatabaseOperation REFRESH = new InsertIdentityOperation(
            DatabaseOperation.REFRESH);

    private static final IColumnFilter DEFAULT_IDENTITY_FILTER = new IColumnFilter()
    {
        public boolean accept(String tableName, Column column)
        {
            return column.getSqlTypeName().endsWith("identity");
        }
    };

    private final DatabaseOperation _operation;

    public InsertIdentityOperation(DatabaseOperation operation)
    {
        _operation = operation;
    }

    private boolean hasIdentityColumn(ITableMetaData metaData,
            IDatabaseConnection connection) throws DataSetException
    {
        DatabaseConfig config = connection.getConfig();
        IColumnFilter identityFilter = ((IColumnFilter) (config
                .getProperty("http://www.dbunit.org/properties/mssql/identityColumnFilter")));
        if (identityFilter == null)
            identityFilter = DEFAULT_IDENTITY_FILTER;
        Column[] columns = metaData.getColumns();
        for (int i = 0; i < columns.length; i++)
        {
            if (identityFilter.accept(metaData.getTableName(), columns[i]))
                return true;
        }
        return false;
    }

    public void execute(IDatabaseConnection connection, IDataSet dataSet)
            throws DatabaseUnitException, SQLException
    {
        Connection jdbcConnection = connection.getConnection();
        Statement statement = jdbcConnection.createStatement();
        try
        {
            IDataSet databaseDataSet = connection.createDataSet();
//            if (!jdbcConnection.getAutoCommit())
//                throw new ExclusiveTransactionException();
//            jdbcConnection.setAutoCommit(false);

            ITableIterator iterator = dataSet.iterator();
            while (iterator.next())
            {
                ITable table = iterator.getTable();
                String tableName = table.getTableMetaData().getTableName();
                ITableMetaData metaData = databaseDataSet
                        .getTableMetaData(tableName);
                boolean hasIdentityColumn = hasIdentityColumn(metaData,
                        connection);
                if (hasIdentityColumn)
                {
                    StringBuffer sqlBuffer = new StringBuffer(128);
                    sqlBuffer.append("SET IDENTITY_INSERT ");
                    sqlBuffer.append(getQualifiedName(connection.getSchema(),
                            metaData.getTableName(), connection));
                    sqlBuffer.append(" ON");
                    statement.execute(sqlBuffer.toString());
                }
                try
                {
                    _operation.execute(connection, new DefaultDataSet(table));
                }
                finally
                {
                    if (hasIdentityColumn)
                    {
                        StringBuffer sqlBuffer = new StringBuffer(128);
                        sqlBuffer.append("SET IDENTITY_INSERT ");
                        sqlBuffer.append(getQualifiedName(connection
                                .getSchema(), metaData.getTableName(),
                                connection));
                        sqlBuffer.append(" OFF");
                        statement.execute(sqlBuffer.toString());
                    }
//                    jdbcConnection.commit();
                }
            }
        }
        finally
        {
//            jdbcConnection.setAutoCommit(true);
            statement.close();
        }
    }
}
